#include "../DWMRICore/DWI.h"
#include "../DWMRICore/DTI.h"

#include <teem/ten.h>

#include <time.h>
#define _USE_MATH_DEFINES
#include <math.h>

char *dwi_pathname = NULL;
char *dti_pathname = NULL;
char *dst_pathname = NULL;

float MLSSize = 2.0f;
CPoint3F *pSamplePoints = NULL;
float *pSamplePointsWeight = NULL;

void PrintUsage()
{
	printf("Usage: MappingFibers -dwi input_dwi_file -dti input_dti_file -dst output_dwi_file\n");
}

int ParseArguments(const int argc, char *argv[])
{
	for (int i = 1; i < argc; ++i) {
		if (strcmp(argv[i], "-dwi") == 0) {
			dwi_pathname = argv[++i];
		} else if (strcmp(argv[i], "-dti") == 0) {
			dti_pathname = argv[++i];
		} else if (strcmp(argv[i], "-dst") == 0) {
			dst_pathname = argv[++i];
		} else {
			return -1;
		}
	}

	if (dwi_pathname == NULL || dti_pathname == NULL || dst_pathname == NULL)
		return -2;

	return 0;
}

void CreateMLSSamplePoints()
{
	SafeDeleteArray(pSamplePoints);
	SafeDeleteArray(pSamplePointsWeight);
	pSamplePoints = new CPoint3F[SAMPLE_POINT_NUM];
	pSamplePointsWeight = new float[SAMPLE_POINT_NUM];

	pSamplePoints[0].m_x = 0.0f;
	pSamplePoints[0].m_y = 0.0f;
	pSamplePoints[0].m_z = 0.0f;
	pSamplePointsWeight[0] = 1.0f;

	clock_t seed = clock();
	srand(seed);

	for (int i = 1; i < SAMPLE_POINT_NUM; ++i) {
		bool flag = false;
		CPoint3F pt;
		while(!flag) {
			float a = (float)rand() / (float)RAND_MAX * 2.0f * (float)M_PI;
			float b = (float)rand() / (float)RAND_MAX * (float)M_PI;
			float r = (float)rand() / (float)RAND_MAX;

			pt.m_x = r * sinf(a) * sinf(b);
			pt.m_y = r * sinf(a) * cosf(b);
			pt.m_z = r * cosf(a);

			flag = true;
			for (int j = 0; j < i; ++j) {
				float dis = Distance(pSamplePoints[j], pt);
				if (dis < 0.25f) {
					flag = false;
					break;
				}
			}
		}

		pSamplePoints[i] = pt;
		pSamplePointsWeight[i] = InnerProduct(pSamplePoints[i], pSamplePoints[i]) * 4.0f;
		pSamplePointsWeight[i] = expf(pSamplePointsWeight[i] / -2.5f);
	}
}

void DWIFilter_MLS(CPoint3F *pos, const float *dti, CDWI *pSrcDWI, float *dwi)
{
	float m[9];
	TEN_T2M(m, dti);

	CPoint3F samples[SAMPLE_POINT_NUM];
	float length[SAMPLE_POINT_NUM];
	float max_length = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		samples[i].m_x = pSamplePoints[i].m_x * m[0] + pSamplePoints[i].m_y * m[1] + pSamplePoints[i].m_z * m[2];
		samples[i].m_y = pSamplePoints[i].m_x * m[3] + pSamplePoints[i].m_y * m[4] + pSamplePoints[i].m_z * m[5];
		samples[i].m_z = pSamplePoints[i].m_x * m[6] + pSamplePoints[i].m_y * m[7] + pSamplePoints[i].m_z * m[8];
		length[i] = sqrtf(InnerProduct(samples[i], samples[i]));
		if (length[i] > max_length)
			max_length = length[i];
	}

	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		samples[i].m_x = samples[i].m_x / max_length * MLSSize + pos->m_x;
		samples[i].m_y = samples[i].m_y / max_length * MLSSize + pos->m_y;
		samples[i].m_z = samples[i].m_z / max_length * MLSSize + pos->m_z;
	}

	uint n = pSrcDWI->GetN();
	float *tempDWI = new float[n];
	float weight = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		pSrcDWI->GetDWI(samples[i], tempDWI);

		for (uint j = 0; j < n; ++j) {
			dwi[j] += pSamplePointsWeight[i] * tempDWI[j];
		}
		weight += pSamplePointsWeight[i];
	}

	for (uint j = 0; j < n; ++j) {
		dwi[j] = dwi[j] / weight;
	}

	delete[] tempDWI;
}

void DWIFilter(CDWI *pSrcDWI, CDTI *pDTI, CDWI *pDstDWI)
{
	uint n = pSrcDWI->GetN();
	uint w = pSrcDWI->GetWidth();
	uint h = pSrcDWI->GetHeight();
	uint d = pSrcDWI->GetDepth();

	pDstDWI->CreateDWI(n, w, h, d);
	printf("======== running DWI Filter ========\n");
	for (uint z = 0; z < d; ++z) {
		for (uint y = 0; y < h; ++y) {
			for (uint x = 0; x < w; ++x) {
				float *dti = new float[7];
				pDTI->GetDTIFast(x, y, z, dti);

				CPoint3F pos = CPoint3F((float)x, (float)y, (float)z);
				float *dwi = new float[n];
				memset(dwi, 0, sizeof(float)*n);
				DWIFilter_MLS(&pos, dti, pSrcDWI, dwi);

				pDstDWI->SetDWI(x, y, z, dwi);

				SafeDeleteArray(dti);
				SafeDeleteArray(dwi);
			}
		}
		float ff = float(z + 1) / (float)d * 100.0f;
		printf("\r%.2f done", ff);
		fflush(stdout);
	}
	printf("\n======== Done. ========\n");
}

int main(int argc, char *argv[])
{
	if (ParseArguments(argc, argv) != 0) {
		PrintUsage();
		return 0;
	}

	// open DWI and DTI file
	CDWI *pSrcDWI = new CDWI();
	pSrcDWI->OpenDWIFile(dwi_pathname);
	CDTI *pDTI = new CDTI();
	pDTI->OpenDTIFile(dti_pathname);

	// create the output DWI
	CDWI *pDstDWI = new CDWI();

	// run DWI filter
	CreateMLSSamplePoints();
	DWIFilter(pSrcDWI, pDTI, pDstDWI);

	// save DWI file
	pDstDWI->SaveDWIFile(dst_pathname);

	SafeDelete(pSrcDWI);
	SafeDelete(pDTI);
	SafeDelete(pDstDWI);

	return 0;
}